# 导入必要的库
import numpy as np
import os
import output_function as of  # 导入自定义的输出函数模块
from scipy.io import loadmat  # 导入scipy库中的loadmat函数，用于加载.mat文件

# 使用示例
# 设置基础路径，根据实际保存路径进行修改
base_path = './data/'  

# 加载初始数据，包括alx, aly, dt等变量
init_data = loadmat(f'{base_path}initial.mat')
alx = init_data['alx'][0][0]
aly = init_data['aly'][0][0]
dt = init_data['dt'][0][0]

# 加载网格数据，包括xx和yy
xx = init_data['x']
yy = init_data['y']
Ra = init_data['Ra'][0][0]  # 加载Ra变量

# 假设我们从out_iter=10开始，直到8000，且每次增加100
# out_iter表示输出的次数，通过控制它来控制输出的时间间隔
for out_iter in range(10, 801, 10):
    # 将输出次数格式化为四位数字符串，不足四位时前面补零
    formatted_iter = f"{out_iter:04d}"

    # 根据基础路径和格式化后的输出次数构建完整的文件名
    filename = f"{base_path}dat{formatted_iter}.mat"

    # 使用loadmat函数加载指定路径的.mat文件，并将其存储在data变量中
    data = loadmat(filename)

    # 从data中直接访问并提取所需的变量，如wi, phii, T0i, Ti, vex, vey等
    wi = data['wi']
    phii = data['phii']
    T0i = data['T0i']
    Ti = data['Ti']
    vex = data['vex']
    vey = data['vey']

    # 设置基础文件夹路径，用于保存生成的图形文件
    base_folder = f'./figure'

    # 调用output_function模块中的diag_n_pcolor_filepath函数，传入相关参数进行绘图或其他处理
    # 具体函数功能需参考output_function模块的文档或源代码
    of.diag_n_pcolor_filepath(Ti, vey, vex, yy, xx, out_iter, Ra, dt, base_folder)

    # 如果输出次数是100的倍数，则打印当前输出次数
    if out_iter % 100 == 0:
        print(out_iter)

    # 以下是可选的其他绘图函数调用示例（当前处于注释状态），可根据需要取消注释并使用
    # of.plot_and_save_n_density_x_section(x, y, n, n_init0_matrix, Ny // 2, out_iter, title_base="delta n_density Distribution", save_folder_base=base_folder)
    # of.plot_and_save_v_x_section(x, y, v, Ny // 2, out_iter, title_base="V Distribution along X-axis", save_folder_base=base_folder)